package com.dynamic.datasource.router;

import com.dynamic.datasource.exception.NoSuchDataSourceException;
import com.dynamic.datasource.factory.WrapDataSource;
import com.dynamic.datasource.provider.DataSourceProvider;
import com.p6spy.engine.spy.P6DataSource;
import io.seata.rm.datasource.DataSourceProxy;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.util.ReflectionUtils;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * @author 王先堂
 * @Description 数据源路由实现
 * @createTime 2022/03/24/ 15:25:00
 */
@Slf4j
@Data
public class DefaultDataSourceRouter extends AbstractDataSourceRouter implements InitializingBean, DisposableBean {

    /**
     * 所有数据源
     */
    private final Map<String, DataSource> allDataSources = new HashMap<>();

    private Boolean p6spy = false;
    private Boolean seata = false;

    @Resource
    private DataSourceProvider dataSourceProvider;

    @Override
    protected DataSource determineDataSource() {
        //选择数据源
        String name = DataSourceContext.get();
        return getDataSource(name);
    }

    @Override
    public void destroy() throws Exception {
        for (Map.Entry<String, DataSource> item : allDataSources.entrySet()) {
            closeDataSource(item.getKey(), item.getValue());
        }
        log.info("close all DataSource");
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        checkEnv();
        Map<String, DataSource> dataSourceMap = dataSourceProvider.load();
        for (Map.Entry<String, DataSource> item : dataSourceMap.entrySet()) {
            addDataSource(item.getKey(), item.getValue());
        }
        log.info("finish add all DataSource");
    }

    /**
     * 配置依赖检查
     */
    private void checkEnv() {
        if (p6spy) {
            try {
                Class.forName("com.p6spy.engine.spy.P6DataSource");
                log.info("DataSource detect P6SPY plugin and enabled it");
            } catch (Exception e) {
                throw new RuntimeException("DataSource enabled P6SPY ,however without p6spy dependency", e);
            }
        }
        if (seata) {
            try {
                Class.forName("io.seata.rm.datasource.DataSourceProxy");
                log.info("DataSource detect ALIBABA SEATA and enabled it");
            } catch (Exception e) {
                throw new RuntimeException("DataSource enabled ALIBABA SEATA,however without seata dependency", e);
            }
        }
    }

    /**
     * 添加数据源
     *
     * @param name       名称
     * @param dataSource 数据源
     */
    public synchronized void addDataSource(String name, DataSource dataSource) {
        DataSource old = allDataSources.put(name, dataSource);
        if (Objects.nonNull(old)) {
            closeDataSource(name, old);
        }
        log.info("add new DataSource,name: {} ", name);
    }

    /**
     * 获取数据源
     *
     * @param name 数据源名称
     * @return
     */
    public DataSource getDataSource(String name) {
        if (allDataSources.containsKey(name)) {
            return allDataSources.get(name);
        }
        throw new NoSuchDataSourceException("cannot find DataSource by name " + name);
    }

    /**
     * 移除数据源
     *
     * @param name 数据源名称
     */
    public synchronized void removeDataSource(String name) {
        if (allDataSources.containsKey(name)) {
            DataSource dataSource = allDataSources.remove(name);
            closeDataSource(name, dataSource);
            log.info("removed DataSource " + name);
        } else {
            log.warn("cannot find DataSource by name " + name + " to remove");
        }
    }

    /**
     * 关闭数据源
     *
     * @param name       数据源名称
     * @param dataSource 数据源
     */
    public void closeDataSource(String name, DataSource dataSource) {
        try {
            if (dataSource instanceof WrapDataSource) {
                ((WrapDataSource) dataSource).close();
            } else {
                if (seata) {
                    if (dataSource instanceof DataSourceProxy) {
                        DataSourceProxy dataSourceProxy = (DataSourceProxy) dataSource;
                        dataSource = dataSourceProxy.getTargetDataSource();
                    }
                }
                if (p6spy) {
                    if (dataSource instanceof P6DataSource) {
                        Field realDataSourceField = P6DataSource.class.getDeclaredField("realDataSource");
                        realDataSourceField.setAccessible(true);
                        dataSource = (DataSource) realDataSourceField.get(dataSource);
                    }
                }
                Method closeMethod = ReflectionUtils.findMethod(dataSource.getClass(), "close");
                if (closeMethod != null) {
                    closeMethod.invoke(dataSource);
                }
            }
        } catch (Exception e) {
            log.warn("closed datasource named [{}] failed", name, e);
        }
    }
}
